package com.tuniu.agents.agent;

import com.tuniu.agents.common.AgentsConstants;
import com.tuniu.agents.message.AgentMessageType;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;

/**
 * @author Noctis Ma
 * @date 2025/2/24 10:56 AM
 */
public abstract class Agent {

    public abstract String getAgentId();

    /**
     * call with chat model
     * @param conversationId conversation id
     * @param message user message
     * @return chat response
     */
    public final String call(String conversationId, String message) {
        return this.call(conversationId, List.of(new UserMessage(message)));
    }

    /**
     *
     * @param conversationId conversation id
     * @param message user message
     * @return
     */
    public final ChatResponse callResponse(String conversationId, String message) {
        return this.callResponse(conversationId, List.of(new UserMessage(message)));
    }

    /**
     * stream call with chat model
     * @param conversationId conversation id
     * @param message user message
     * @return streaming chat response
     */
    public final Flux<String> stream(String conversationId, String message) {
        return this.stream(conversationId, List.of(new UserMessage(message)));
    }

    /**
     *
     * @param conversationId conversation id
     * @param message user message
     * @return
     */
    public final Flux<ChatResponse> streamResponse(String conversationId, String message) {
        return this.streamResponse(conversationId, List.of(new UserMessage(message)));
    }

    /**
     * call with chat model
     * @param conversationId conversation id
     * @param messages user message
     * @return chat response
     */
    public final String call(String conversationId, List<Message> messages) {
        return call(conversationId, messages, null);
    }

    /**
     *
     * @param conversationId conversation id
     * @param messages user message
     * @return
     */
    public final ChatResponse callResponse(String conversationId, List<Message> messages) {
        return callResponse(conversationId, messages, null);
    }

    /**
     * stream call with chat model
     * @param conversationId conversation id
     * @param messages user message
     * @return streaming chat response
     */
    public final Flux<String> stream(String conversationId, List<Message> messages) {
        return this.stream(conversationId, messages, null);
    }

    /**
     * stream call with chat model
     * @param conversationId conversation id
     * @param messages user message
     * @return
     */
    public final Flux<ChatResponse> streamResponse(String conversationId, List<Message> messages) {
        return this.streamResponse(conversationId, messages, null);
    }

    /**
     * stream call with chat model
     * @param conversationId conversation id
     * @param messages user message
     * @param context context
     * @return
     */
    public abstract Flux<ChatResponse> streamResponse(String conversationId, List<Message> messages, Map<String,Object> context);

    /**
     * stream call with chat model
     * @param conversationId conversation id
     * @param messages user message
     * @param context context
     * @return
     */
    public abstract Flux<String> stream(String conversationId, List<Message> messages, Map<String,Object> context);

    /**
     * call with chat model
     * @param conversationId conversation id
     * @param messages user message
     * @param context context
     * @return chat response
     */
    public abstract ChatResponse callResponse(String conversationId, List<Message> messages, Map<String,Object> context);

    /**
     * call with chat model
     * @param conversationId conversation id
     * @param messages user message
     * @param context context
     * @return chat response
     */
    public abstract String call(String conversationId, List<Message> messages, Map<String,Object> context);

    /**
     * build chatClientRequestSpec
     * @param conversationId
     * @param messages
     * @return
     */
    public abstract ChatClient.ChatClientRequestSpec chatClientRequestSpec(String conversationId, List<Message> messages,Map<String,Object> context);

    public ChatResponse wrapper(ChatResponse chatResponse, AgentMessageType messageType){
        if(chatResponse.getMetadata().containsKey(AgentsConstants.AgentsConversationEnvironmentKey.AGENT_ID)){
            return chatResponse;
        } else {
            return ChatResponse.builder().from(chatResponse)
                    .metadata(AgentsConstants.AgentResponseMetadataKey.METADATA_AGENT_ID, this.getAgentId())
                    .metadata(AgentsConstants.AgentResponseMetadataKey.METADATA_MESSAGE_TYPE_KEY, messageType)
                    .build();
        }
    }

    public ChatResponse wrapper(ChatResponse chatResponse){
        return this.wrapper(chatResponse, AgentMessageType.LLM_TEXT);
    }
}
